//  BitWrk - A Bitcoin-friendly, anonymous marketplace for computing power
//  Copyright (C) 2013-2019 Jonas Eschenburg <jonas@bitwrk.net>
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <http://www.gnu.org/licenses/>.

package remotesync

import (
	"errors"
	"fmt"
	"github.com/indyjo/cafs"
	"github.com/indyjo/cafs/remotesync/shuffle"
	"io"
	"log"
)

// By passing a callback function to some of the transmissions functions,
// the caller may subscribe to the current transmission status.
type TransferStatusCallback func(bytesToTransfer, bytesTransferred int64)

// Interface Chunks allows iterating over any sequence of chunks.
type Chunks interface {
	// Function NextChunk returns either of three cases:
	// - A File, nil   (good case)
	// - nil, io.EOF   (terminal case: end of stream)
	// - nil, an error (terminal case: an error occurred)
	// It is the caller's duty to call Dispose() on the file returned.
	NextChunk() (cafs.File, error)

	// Function Dispose must be called when this object is no longer used.
	Dispose()
}

// Function ChunksOfFiles returns the chunks of a File as an implementation of the Chunks
// interface. It's the caller's responsibility to call Dispose() on the returned object.
func ChunksOfFile(file cafs.File) Chunks {
	return chunksOfFile{iter: file.Chunks()}
}

// Struct chunksOfFile is a minimal wrapper around a FileIterator that implements
// the Chunks interface.
type chunksOfFile struct {
	iter cafs.FileIterator
}

func (c chunksOfFile) NextChunk() (cafs.File, error) {
	if c.iter.Next() {
		return c.iter.File(), nil
	}
	return nil, io.EOF
}

func (c chunksOfFile) Dispose() {
	c.iter.Dispose()
}

// Iterates over a wishlist (read from `r` and pertaining to a permuted order of hashes),
// and calls `f` for each chunk of `file`, requested or not.
// If `f` returns an error, aborts the iteration and also returns the error.
func forEachChunk(chunks Chunks, r io.ByteReader, perm shuffle.Permutation, f func(chunk cafs.File, requested bool) error) error {
	bits := newBitReader(r)

	// Prepare shuffler for iterating the file's chunks in shuffled order, matching them with
	// whishlist bits and calling `f` for each chunk, requested or not.
	shuffler := shuffle.NewStreamShuffler(perm, nil, func(v interface{}) error {
		var requested bool
		if b, err := bits.ReadBit(); err != nil {
			return fmt.Errorf("error reading from wishlist bitstream: %v", err)
		} else {
			requested = b
		}

		if v == nil {
			// This is a placeholder key generated by the shuffler. Require that the receiver
			// signalled not to request the corresponding chunk.
			if requested {
				return errors.New("receiver requested the empty chunk")
			}
			// otherwise, there's nothing to do
			return nil
		}

		// We have a chunk with a corresponding wishlist bit. Dispatch to delegate function.
		chunk := v.(cafs.File)
		err := f(chunk, requested)
		chunk.Dispose()
		return err
	})

	// At the end of this function, we must make sure that all chunks still stored
	// in the shuffler are disposed of.
	defer func() {
		s := shuffler.WithFunc(func(v interface{}) error {
			if v != nil {
				v.(cafs.File).Dispose()
			}
			return nil
		})
		_ = s.End()
	}()

	// Iterate through the chunks and put their keys into the shuffler.
	for {
		if chunk, err := chunks.NextChunk(); err == nil {
			if err := shuffler.Put(chunk); err != nil {
				return err
			}
		} else if err == io.EOF {
			break
		} else {
			return err
		}

	}
	if err := shuffler.End(); err != nil {
		return err
	}

	// Expect whishlist byte stream to be read completely
	if _, err := r.ReadByte(); err != io.EOF {
		return errors.New("wishlist too long")
	}
	return nil
}

// Writes a stream of chunk length / data pairs, permuted by a shuffler corresponding to `perm`,
// into an io.Writer, based on the chunks of a file and a matching permuted wishlist of requested chunks,
// read from `r`.
func WriteChunkData(chunks Chunks, bytesToTransfer int64, r io.ByteReader, perm shuffle.Permutation, w FlushWriter, cb TransferStatusCallback) error {
	if LoggingEnabled {
		log.Printf("Sender: Begin WriteChunkData")
		defer log.Printf("Sender: End WriteChunkData")
	}

	// Determine the number of bytes to transmit by starting at the maximum and subtracting chunk
	// total whenever we read a 0 (chunk not requested)
	if cb != nil {
		cb(bytesToTransfer, 0)
	}

	// Iterate requested chunks. Write the chunk's length (as varint) and the chunk data
	// into the output writer. Update the number of bytes transferred on the go.
	var bytesTransferred int64
	return forEachChunk(chunks, r, perm, func(chunk cafs.File, requested bool) error {
		if requested {
			if err := writeVarint(w, chunk.Size()); err != nil {
				return err
			}
			r := chunk.Open()
			if n, err := io.Copy(w, r); err != nil {
				_ = r.Close()
				return err
			} else {
				w.Flush()
				bytesTransferred += n
			}
			if err := r.Close(); err != nil {
				return err
			}
		} else {
			bytesToTransfer -= chunk.Size()
		}
		if cb != nil {
			// Notify callback of status
			cb(bytesToTransfer, bytesTransferred)
		}
		return nil
	})
}
